
#include "conjgrad_mt_state.h"
#include "linmin.h"

#include <gsl/gsl_multimin.h>
#include <gsl/gsl_blas.h>

static int conjgrad_mt_alloc(void *vstate, size_t n)
{
  conjgrad_mt_state_t *state = (conjgrad_mt_state_t *)vstate;

  state->line_minimizer = lineminimizer_alloc(lineminimizer_morethuente, n);
  state->d = gsl_vector_alloc(n);
  state->gradient_plus = gsl_vector_alloc(n);
  state->x_plus = gsl_vector_alloc(n);
  
  return GSL_SUCCESS;
}

static int conjgrad_mt_set(void *vstate, gsl_multimin_function_fdf * fdf,
                           const gsl_vector * x, double *f, gsl_vector * gradient,
                           double step_size, double tol)
{
  conjgrad_mt_state_t *state = (conjgrad_mt_state_t *) vstate;
  
  state->alpha0 = step_size;
  state->tol = tol;
  state->ci = 0;
  state->delta_f = 0.0;
  
  GSL_MULTIMIN_FN_EVAL_F_DF (fdf, x, f, gradient);
	
  lineminimizer_morethuente_set_gamma(state->line_minimizer, 0.0);

  gsl_vector_memcpy (state->d, gradient);
  gsl_blas_dscal(-1.0, state->d);
  
  return GSL_SUCCESS;
}

static void conjgrad_mt_free (void *vstate)
{
  conjgrad_mt_state_t *state = (conjgrad_mt_state_t *) vstate;

  gsl_vector_free (state->x_plus);
  gsl_vector_free (state->gradient_plus);
  gsl_vector_free (state->d);
  lineminimizer_free(state->line_minimizer);
}

static int conjgrad_fr_mt_iterate(void *vstate, gsl_multimin_function_fdf *fdf,
                                  gsl_vector * x, double *f,
                                  gsl_vector * gradient, gsl_vector * dx)
{
  double alpha, alpha0;
  double dphi0;
  int status;
  double gnorm_sq;
  double gnorm_plus_sq;
  double gamma;
  
  conjgrad_mt_state_t *state = (conjgrad_mt_state_t *) vstate;
  
  gsl_blas_ddot(gradient, state->d, &dphi0);
  /* reset back to steepest descent direction if the */
  /* current search direction is not a descent direction */
  if(dphi0 >= 0.0)
    {
      gsl_blas_dcopy(gradient, state->d);
      gsl_blas_dscal(-1.0, state->d);
      gsl_blas_ddot(gradient, gradient, &dphi0);
      dphi0 = -dphi0;
    }
  /* compute the initial step length with Fletcher's method */
  alpha0 = lineminimizer_initstep_fletcher(dphi0, state->delta_f);
  
  /* obtain step length */
  status = lineminimizer_minimize(state->line_minimizer, fdf, alpha0, state->tol,
                                  *f, x, gradient, state->d, &alpha, 
                                  state->x_plus, &state->f_plus, state->gradient_plus);
  
  if(gsl_isinf(state->f_plus) || gsl_isnan(state->f_plus))
    return GSL_ERUNAWAY;
  
  /* compute gamma = g'.g' / g.g */
  gsl_blas_ddot(state->gradient_plus, state->gradient_plus, &gnorm_plus_sq);
  gsl_blas_ddot(gradient, gradient, &gnorm_sq);
  gamma = gnorm_plus_sq / gnorm_sq;
    
  /* restart every n+1 iterations */
  if(state->ci == fdf->n)
    {
      gamma = 0.0;
      state->ci = 0;
    }
  
  lineminimizer_morethuente_set_gamma(state->line_minimizer, 0.0);
	
  /* update the search direction by applying */
  /* d' = -g + gamma * d */
  gsl_blas_dscal(gamma, state->d);
  gsl_blas_daxpy(-1.0, state->gradient_plus, state->d);
  
  state->delta_f = *f - state->f_plus;
  
  /* update the current iterate and its function and gradient values */
  gsl_blas_dcopy(state->x_plus, dx);
  gsl_blas_daxpy(-1.0, x, dx);
  gsl_blas_dcopy(state->x_plus, x);
  gsl_blas_dcopy(state->gradient_plus, gradient);
  *f = state->f_plus;
  
  state->ci++;
  
  return GSL_SUCCESS;
}

static int conjgrad_pr_mt_iterate(void *vstate, gsl_multimin_function_fdf * fdf,
                                  gsl_vector * x, double *f,
                                  gsl_vector * gradient, gsl_vector * dx)
{
  double alpha, alpha0;
  double dphi0;
  int status;
  double gnorm_sq;
  double gnorm_plus_sq;
  double gpg;
  double gamma;
  
  conjgrad_mt_state_t *state = (conjgrad_mt_state_t *) vstate;
  
  gsl_blas_ddot(gradient, state->d, &dphi0);
  /* reset back to steepest descent direction if the */
  /* current search direction is not a descent direction */
  if(dphi0 >= 0.0)
  {
    gsl_blas_dcopy(gradient, state->d);
    gsl_blas_dscal(-1.0, state->d);
    gsl_blas_ddot(gradient, gradient, &dphi0);
    dphi0 = -dphi0;
  }
  /* compute initial step length with Fletcher's method */
  alpha0 = lineminimizer_initstep_fletcher(dphi0, state->delta_f);
  
  /* obtain step length */
  status = lineminimizer_minimize(state->line_minimizer, fdf, alpha0, state->tol,
                                  *f, x, gradient, state->d, &alpha,
                                  state->x_plus, &state->f_plus, state->gradient_plus);
  
  if(gsl_isinf(state->f_plus) || gsl_isnan(state->f_plus))
    return GSL_ERUNAWAY;
  
  /* update the search direction by applying */
  /* d' = -g + gamma * d */
  gsl_blas_ddot(gradient, gradient, &gnorm_sq);
  gsl_blas_ddot(state->gradient_plus, state->gradient_plus, &gnorm_plus_sq);
  gsl_blas_ddot(state->gradient_plus, gradient, &gpg);
    
  /* compute gamma = (g'-g).g' / g.g */
  /* and ensure that it is nonnegative */
  gamma = GSL_MAX_DBL((gnorm_plus_sq - gpg) / gnorm_sq, 0.0);
	
  lineminimizer_morethuente_set_gamma(state->line_minimizer, gamma);
  
  gsl_blas_dscal(gamma, state->d);
  gsl_blas_daxpy(-1.0, state->gradient_plus, state->d);
  
  state->delta_f = *f - state->f_plus;
  
  /* update the current iterate and its function and gradient values */
  gsl_blas_dcopy(state->x_plus, dx);
  gsl_blas_daxpy(-1.0, x, dx);
  gsl_blas_dcopy(state->x_plus, x);
  gsl_blas_dcopy(state->gradient_plus, gradient);
  *f = state->f_plus;
  
  return GSL_SUCCESS;
}

static int conjgrad_mt_restart (void *vstate)
{
  /* reset back to steepest descent direction */
  conjgrad_mt_state_t *state = (conjgrad_mt_state_t *) vstate;
  
  gsl_blas_dcopy(state->gradient_plus, state->d);
  gsl_blas_dscal(-1.0, state->d);
  
  return GSL_SUCCESS;
}

static const gsl_multimin_fdfminimizer_type conjgrad_fr_mt_type = 
{
  "conjgrad_fr_mt", 
  sizeof (conjgrad_mt_state_t),
  &conjgrad_mt_alloc,
  &conjgrad_mt_set,
  &conjgrad_fr_mt_iterate, 
  &conjgrad_mt_restart, 
  &conjgrad_mt_free
};

const gsl_multimin_fdfminimizer_type
* gsl_ext_multimin_fdfminimizer_conjgrad_fr_mt = &conjgrad_fr_mt_type;

static const gsl_multimin_fdfminimizer_type conjgrad_pr_mt_type = 
{
  "conjgrad_pr_mt",
  sizeof (conjgrad_mt_state_t),
  &conjgrad_mt_alloc,
  &conjgrad_mt_set,
  &conjgrad_pr_mt_iterate,
  &conjgrad_mt_restart,
  &conjgrad_mt_free
};

const gsl_multimin_fdfminimizer_type
* gsl_ext_multimin_fdfminimizer_conjgrad_pr_mt = &conjgrad_pr_mt_type;
